import torch
from ddl_platform.ddlib import dopt

rank = dopt.rank()
torch.cuda.set_device(rank)
dopt.init()

n=4
k=2
t1 = torch.rand(n).cuda()
_, indexes = torch.topk(t1, k=k)
values = t1[indexes]
sparse_tensor = torch.zeros_like(t1)
print('before: ', rank, t1, values, indexes)
send_values, send_indexes, excluded_indexes = dopt.gtopk_sparse_recursive_allreduce(dopt.comm, sparse_tensor, values, indexes)
print('after: ', rank, send_values, send_indexes, excluded_indexes)
